import numpy as np
from scipy.stats import kde
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats
#from mpl_toolkits.mplot3d import Axes3D 
from math import factorial
from itertools import chain
from matplotlib import rcParams
import pandas as pd
import random

rad2deg = 180/np.pi
def readrmc6f(filename):
    lattice = np.zeros((3,3))
    with open (filename) as f:
        clines = f.readlines()
        f.close()
    idx = clines.index('Atoms:\n')
    metadata = clines[:idx+1]
    for line in metadata:
        if line.find('Number of types of atoms') != -1:
            ntypes = int(line.rstrip().split(':')[1])
        if line.find('Atom types present') != -1:
            types = line.rstrip().split(':')[1].split()
        if line.find('Number of atoms:') != -1:
            natoms = int(line.rstrip().split(':')[1])
        if line.find('Number of each atom type:') != -1:
            typenatoms = line.rstrip().split(':')[1].split() # note typeatoms is a string
    for i in range (3):
        for j in range(3):
            lattice[i,j] = float(metadata[idx+i-3].split()[j])
    #lattice['x']=float(metadata[idx-3].split()[0])
    #lattice['y']=float(metadata[idx-2].split()[1])
    #lattice['z']=float(metadata[idx-1].rstrip().split()[2])
    #there arefloat( two lattice parameters in the rmc6f. I recall the the matrix value are the values used to 
    #calculate the atomic coordinates but that's not sure for reduced.rmc6f
    
    atomtext=clines[idx+1:]
    if len(atomtext) != natoms:
        print ('atomtext does not include the right number of atoms')
    atoms={}
    for line in atomtext:
        iatom = int(line.split()[0])   
        atoms[iatom]=[None] * 9
        atoms[iatom][0] = line.split()[1]   #atomtype
        if line.split()[2][2].isdigit():
            atoms[iatom][1] = int(line.split()[2][1:3])
        else:
            atoms[iatom][1] = int(line.split()[2][1:2])                 #index within atom type  
        atoms[iatom][2] = float(line.split()[3])  #x  
        atoms[iatom][3] = float(line.split()[4])  #y
        atoms[iatom][4] = float(line.split()[5])  #z 
        atoms[iatom][5] = int(line.split()[6])              #atom index within original unitcell?
        atoms[iatom][6] = int(line.split()[7])              #cell_reference a
        atoms[iatom][7] = int(line.split()[8])              #cell_reference b
        atoms[iatom][8] = int(line.split()[9])              #cell_reference c
            
    return atoms, lattice, metadata

def rewritermc6f_atomtype(filename,atoms,atomtype):
    f = open(filename,'w+')
   # print(len(atoms))
    for iatom in range(1,len(atoms)+1):
        #print(atoms[iatom][1],atomtype[atoms[iatom][1]])
        f.write('{} {} [{}] {:.10f} {:<.10f} {:<.10f} {} {} {} {}\n'.format(iatom, atomtype[atoms[iatom][1]],atoms[iatom][1],atoms[iatom][2],
            atoms[iatom][3],atoms[iatom][4],atoms[iatom][5],atoms[iatom][6],atoms[iatom][7],atoms[iatom][8]))
    f.close()

def rewritermc6f(filename,atoms):
    f = open(filename,'w+')
   # print(len(atoms))
    f.writelines (metadata)
    for iatom in range(1,len(atoms)+1):
        #print(atoms[iatom][1],atomtype[atoms[iatom][1]])
        f.write('{} {} [{}] {:.10f} {:<.10f} {:<.10f} {} {} {} {}\n'.format(iatom, atoms[iatom][0],atoms[iatom][1],atoms[iatom][2],
            atoms[iatom][3],atoms[iatom][4],atoms[iatom][5],atoms[iatom][6],atoms[iatom][7],atoms[iatom][8]))
    f.close()

def portion_select(arr, portion):
    # 计算选择的元素数量
    sample_size = int(len(arr) * portion)
    # 随机选择30%的元素
    return random.sample(arr, sample_size)

def distances(xyz):#xyz对应于存储xyz坐标的list，e.g.coorbondsvectors[i][j][k]
    dis = round(np.sqrt(xyz[0]**2+xyz[1]**2+xyz[2]**2),4)
    return dis
def findneighbours(atoms,lattice,centre,r,dr):
    natoms = len(atoms)
    rmin = r-dr
    rmax = r+dr
    icentre = 0
    coor =[]
    for i in range (natoms):
        iatom = i+1
        if atoms[iatom][0] == centre:
            coor.append(icentre)
            coor[icentre]=[]
            coor[icentre].append(iatom)       
            icentre = icentre + 1
    for i in range (icentre):
        iatom = coor[i][0]
        for j in range (natoms):
            jatom = j+1
            dx  = atoms[iatom][2] - atoms[jatom][2] + 1.5
            dx  = dx - int(dx) - 0.5
            dy  = atoms[iatom][3] - atoms[jatom][3] + 1.5
            dy  = dy - int(dy) - 0.5
            dz  = atoms[iatom][4] - atoms[jatom][4] + 1.5
            dz  = dz - int(dz) - 0.5
            adx = lattice[0][0]*dx + lattice[0][1]*dy + lattice[0][2]*dz
            ady = lattice[1][0]*dx + lattice[1][1]*dy + lattice[1][2]*dz
            adz = lattice[2][0]*dx + lattice[2][1]*dy + lattice[2][2]*dz
            ar = np.sqrt(adx**2+ady**2+adz**2)
            if rmin<=ar<=rmax:
                    coor[i].append(jatom)
    return coor

def neighbours(rmc6f,centre,r,dr):
    atoms,lattice = readrmc6f(rmc6f)
    natoms = len(atoms)
    rmin = r-dr
    rmax = r+dr
    icentre = 0
    coor =[]
    for i in range (natoms):
        iatom = i+1
        if atoms[iatom][0] == centre:
            coor.append(icentre)
            coor[icentre]=[]
            coor[icentre].append(iatom)       
            icentre = icentre + 1
    for i in range (icentre):
        iatom = coor[i][0]
        for j in range (natoms):
            jatom = j+1
            dx  = atoms[iatom][2] - atoms[jatom][2] + 1.5
            dx  = dx - int(dx) - 0.5
            dy  = atoms[iatom][3] - atoms[jatom][3] + 1.5
            dy  = dy - int(dy) - 0.5
            dz  = atoms[iatom][4] - atoms[jatom][4] + 1.5
            dz  = dz - int(dz) - 0.5
            adx = lattice[0][0]*dx + lattice[0][1]*dy + lattice[0][2]*dz
            ady = lattice[1][0]*dx + lattice[1][1]*dy + lattice[1][2]*dz
            adz = lattice[2][0]*dx + lattice[2][1]*dy + lattice[2][2]*dz
            ar = np.sqrt(adx**2+ady**2+adz**2)
            if rmin<=ar<=rmax:
                    coor[i].append(jatom)
    return coor

def countneighbours(coor):
    NCoor=[]
    for i in range (len(coor)):
        NCoor.append(len(coor[i])-1)
    return NCoor

def splitcoor(atoms,coor,sites):
    sitecoor =[]
    for i in range (len(sites)):
        sitecoor.append([])
    for i in range (len(coor)):
        iatom = coor[i][0]
        index = atoms[iatom][1]
        for j in range (len(sites)):
            if index == sites[j]:
                sitecoor[j].append(coor[i])                    
    return sitecoor
 
def calculate_coorbondsvectors(atoms,lattice,sitecoor):
    coorbondsvectors =[]
    for i in range (len(sitecoor)): # i 对应 第几种中心配位原子
        coorbondsvectors.append([])  
        for j in range (len(sitecoor[i])): #j对应确定index种类后的第几个配位中心原子
            iatom = sitecoor[i][j][0]
            coorbondsvectors[i].append([])
            for k in range(len(sitecoor[i][j])-1):
                coorbondsvectors[i][j].append([])
                jatom = sitecoor[i][j][k+1]
                dx  = atoms[iatom][2] - atoms[jatom][2] + 1.5
                dx  = dx - int(dx) - 0.5
                dy  = atoms[iatom][3] - atoms[jatom][3] + 1.5
                dy  = dy - int(dy) - 0.5
                dz  = atoms[iatom][4] - atoms[jatom][4] + 1.5
                dz  = dz - int(dz) - 0.5
                adx = lattice[0][0]*dx + lattice[0][1]*dy + lattice[0][2]*dz
                ady = lattice[1][0]*dx + lattice[1][1]*dy + lattice[1][2]*dz
                adz = lattice[2][0]*dx + lattice[2][1]*dy + lattice[2][2]*dz
                #ar = np.sqrt(adx**2+ady**2+adz**2)
                #转成正交坐标系的bond vector坐标储存在 coorbondsvectors[i][j][k]里
                #coorbondsvectors[i][j]对应第几类第几个中心原子，去对应与一个list consist of bond vectors
                coorbondsvectors[i][j][k].append(float(adx))#x对应coorbondsvectors[i][j][k][0]
                coorbondsvectors[i][j][k].append(float(ady))#y对应coorbondsvectors[i][j][k][1]
                coorbondsvectors[i][j][k].append(float(adz))#z对应coorbondsvectors[i][j][k][2]
    return coorbondsvectors

def unit_vector(vector):
    """ Returns the unit vector of the vector.  """
    return vector / np.linalg.norm(vector)

def angle_between(v1, v2):
    """ Returns the angle in radians between vectors 'v1' and 'v2'::

            >>> angle_between((1, 0, 0), (0, 1, 0))
            1.5707963267948966
            >>> angle_between((1, 0, 0), (1, 0, 0))
            0.0
            >>> angle_between((1, 0, 0), (-1, 0, 0))
            3.141592653589793
    """
    v1_u = unit_vector(v1)
    v2_u = unit_vector(v2)
    radians = np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))
    degree = round(180*radians/np.pi,4)
    
    return degree

def calculate_angles(bondvector):# this calculate all angles in a list of bond vectors and return a list of angles
    angles = []
    for i in range(len(bondvector)):
        for k in range(i+1,len(bondvector)):
            angles.append(angle_between(bondvector[i],bondvector[k]))
    return angles   

def distances(xyz):#xyz对应于存储xyz坐标的list，e.g.coorbondsvectors[i][j][k]
    dis = round(np.sqrt(xyz[0]**2+xyz[1]**2+xyz[2]**2),4)
    return dis

def calculate_distances(bondvector):
    dis=[]
    for i in range(len(bondvector)):
        dis.append(distances(bondvector[i]))
    return dis

def calculate_bonds(coorbondsvectors):
    bonds = []
    for i in range(len(coorbondsvectors)):
        bonds.append([])
        for j in range (len(coorbondsvectors[i])):
            bonds[i].append(calculate_distances(coorbondsvectors[i][j]))
    #bonds[i][j][k]对应第i类第j个中心原子的第k个配位原子距离,iatom,jatom序号可从sitecoor[i][j][0],sitecoor[i][j][k+1]查
    return bonds

def bondshist(bonds):
    bondhist=[]
    for i in range(len(bonds)):
        for j in range(len(bonds[i])):
            for k in range(len(bonds[i][j])):
                bondhist.append([bonds[i][j][k],[i,j,k]])    
    return bondhist

def choosebonds(bondhist,rmin=0,rmax = 5):
    choosed_bonds =[]
    centres=[]
    r_range = range(int(rmin*10000),int(rmax*10000)+1)
    for i in range(len(bondhist)):
        if int(bondhist[i][0]*10000) in r_range:
            #print (bondhist[i][0])
            choosed_bonds.append(bondhist[i])
            centres.append(bondhist[i][1][0]) 
    return choosed_bonds,centres

def find_bonds(atoms,lattice,sitecoor,centre_index,ligand_index):
    coorbonds =[]
    for i in range (len(sitecoor)): # i 对应 第几种中心配位原子
        iatom = sitecoor[i][0][0]
        if int(atoms[iatom][5]) in centre_index:
            for j in range (len(sitecoor[i])): #j对应确定index种类后的第几个配位中心原子
                iatom = sitecoor[i][j][0]
                for k in range(len(sitecoor[i][j])-1):
                    jatom = sitecoor[i][j][k+1]
                    if int(atoms[jatom][5]) in ligand_index:
                        dx  = atoms[iatom][2] - atoms[jatom][2] + 1.5
                        dx  = dx - int(dx) - 0.5
                        dy  = atoms[iatom][3] - atoms[jatom][3] + 1.5
                        dy  = dy - int(dy) - 0.5
                        dz  = atoms[iatom][4] - atoms[jatom][4] + 1.5
                        dz  = dz - int(dz) - 0.5
                        adx = lattice[0][0]*dx + lattice[0][1]*dy + lattice[0][2]*dz
                        ady = lattice[1][0]*dx + lattice[1][1]*dy + lattice[1][2]*dz
                        adz = lattice[2][0]*dx + lattice[2][1]*dy + lattice[2][2]*dz
                        coorbonds.append([adx,ady,adz])
    return coorbonds

def find_bonds_dis(coorbonds):
    bonds_dis=[]
    for i in range(len(coorbonds)):
        bonds_dis.append(distances(coorbonds[i]))
    return bonds_dis